package com.xuhui.xuaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;


import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;


/**
 * 基于文件持久化的对话记忆
 * @author 21829
 */
public class FileBasedChatMemory implements ChatMemory {

    private   final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File file = new File(dir);
        if (!file.exists()) {
            file.mkdirs();
        }
    }


    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> chatMemory4FileMessages = getChatMemory4File(conversationId);
        chatMemory4FileMessages.addAll(messages);
        saveChatMemory2File(conversationId, chatMemory4FileMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messages = getChatMemory4File(conversationId);
        return messages.stream()
                .skip(Math.max(0,messages.size() - lastN))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File file = getChatMemoryFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    private List<Message> getChatMemory4File(String conversationId) {
        File file = getChatMemoryFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (file.exists()) {
            try(Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (FileNotFoundException e) {
                e.printStackTrace();
            }
        }
        return messages;
    }



    public void saveChatMemory2File(String conversationId, List<Message> messages) {
        File file = getChatMemoryFile(conversationId);
        try(Output output = new Output(new FileOutputStream(file))){
           kryo.writeObject(output, messages);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        }
    }



    public File getChatMemoryFile(String conversationId) {
        return new File(BASE_DIR, conversationId+".kryo");
    }
}
